extern crate fish_utils as fu;
use fu::graph::Graph;
use fu::hypergraph::HyperGraph;
use numpy::{IntoPyArray, PyArray1, PyArray2, PyArray3, PyReadonlyArray1, PyReadonlyArray2};
use pyo3::prelude::*;

#[pyclass(name = "Graph")]
struct PyGraph {
    g: Graph<f32>,
}

#[pymethods]
impl PyGraph {
    #[new]
    fn new(
        data: PyReadonlyArray1<f32>,
        row_indices: PyReadonlyArray1<usize>,
        col_indices: PyReadonlyArray1<usize>,
        num_nodes: usize,
    ) -> Self {
        let g = Graph::<f32>::new(
            data.as_array(),
            row_indices.as_array(),
            col_indices.as_array(),
            num_nodes,
        );
        PyGraph { g: g }
    }

    fn num_node(&self) -> usize {
        self.g.num_node()
    }

    fn num_edge(&self) -> usize {
        self.g.num_edge()
    }

    fn random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<'py, usize>,
        length: usize,
        num_path: usize,
    ) -> PyResult<Bound<'py, PyArray3<usize>>> {
        Ok(self
            .g
            .random_walk(&nodes.as_array(), length, num_path)
            .into_pyarray_bound(py))
    }

    fn k_hop<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        k: usize,
    ) -> PyResult<Bound<'py, PyArray1<usize>>> {
        Ok(self.g.k_hop(&nodes.as_array(), k).into_pyarray_bound(py))
    }

    fn restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rate: f32,
    ) -> PyResult<Bound<'py, PyArray3<usize>>> {
        Ok(self
            .g
            .restart_random_walk(&nodes.as_array(), length, num_path, rate)
            .into_pyarray_bound(py))
    }

    fn sample_subgraph_for_bpr<'py> (&self, py: Python<'py>, edges: PyReadonlyArray2<usize>, length: usize, num_path: usize, rwr_rate: f32, num_sample: usize) -> PyResult<(
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray2<usize>>,
    )> {
        let (a,b,c,d,e) = self.g.sample_subgraph_for_bpr(&edges.as_array(), length, num_path, rwr_rate, num_sample);
        Ok((a.into_pyarray_bound(py),b.into_pyarray_bound(py),c.into_pyarray_bound(py), d.into_pyarray_bound(py), e.into_pyarray_bound(py)))
    }

    fn sample_subgraph_from_edges<'py> (&self, py: Python<'py>, edges: PyReadonlyArray2<usize>, length: usize, num_path: usize, rwr_rate: f32) -> PyResult<(
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
    )> {
        let (a,b,c,d) = self.g.sample_subgraph_from_edges_by_restart_random_walk(&edges.as_array(), length, num_path, rwr_rate);

        Ok((a.into_pyarray_bound(py),b.into_pyarray_bound(py),c.into_pyarray_bound(py),  d.into_pyarray_bound(py)))
    }

    fn sample_linkprediction_by_restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        rate: f32,
        num_sample: usize,
    ) -> PyResult<(
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<f32>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray2<usize>>,
        Bound<'py, PyArray2<usize>>,
    )> {
        let (a, b, c, d, e) = self.g.sample_linkprediction_by_restart_random_walk(
            &nodes.as_array(),
            length,
            num_path,
            rwr_rate,
            rate,
            num_sample,
        );
        let (ba, bb, bc) = b.tocoo().to_ndarray();
        Ok((
            a.into_pyarray_bound(py),
            ba.into_pyarray_bound(py),
            bb.into_pyarray_bound(py),
            bc.into_pyarray_bound(py),
            c.into_pyarray_bound(py),
            d.into_pyarray_bound(py),
            e.into_pyarray_bound(py),
        ))
    }

    fn sample_by_restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
    ) -> PyResult<(
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<f32>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray1<usize>>,
    )> {
        let (a, b, c) =
            self.g
                .sample_by_restart_random_walk(&nodes.as_array(), length, num_path, rwr_rate);
        let (ba, bb, bc) = b.tocoo().to_ndarray();
        Ok((
            a.into_pyarray_bound(py),
            ba.into_pyarray_bound(py),
            bb.into_pyarray_bound(py),
            bc.into_pyarray_bound(py),
            c.into_pyarray_bound(py),
        ))
    }
}

#[pyclass(name = "HyperGraph")]
struct PyHyperGraph {
    hg: HyperGraph<bool>,
}

#[pymethods]
impl PyHyperGraph {
    #[new]
    fn new(
        data: PyReadonlyArray1<bool>,
        row_indices: PyReadonlyArray1<usize>,
        col_indices: PyReadonlyArray1<usize>,
        num_nodes: usize,
        num_edges: usize,
    ) -> Self {
        let hg = HyperGraph::<bool>::new(
            data.as_array(),
            row_indices.as_array(),
            col_indices.as_array(),
            num_nodes,
            num_edges,
        );
        PyHyperGraph { hg: hg }
    }

    fn num_node(&self) -> usize {
        self.hg.num_node()
    }

    fn num_edge(&self) -> usize {
        self.hg.num_edge()
    }

    fn restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rate: f32,
    ) -> PyResult<Bound<'py, PyArray3<usize>>> {
        Ok(self
            .hg
            .restart_random_walk(&nodes.as_array(), length, num_path, rate)
            .into_pyarray_bound(py))
    }

    fn sample_linkprediction_by_restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        rate: f32,
        num_e: usize,
        num_n: usize,
    ) -> PyResult<(
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyArray2<usize>>,
        Bound<'py, PyArray3<usize>>,
        Bound<'py, PyArray3<usize>>,
        Bound<'py, PyArray3<usize>>,
    )> {
        let (a, b, c, d, e) = self.hg.sample_linkprediction_by_restart_random_walk(
            &nodes.as_array(),
            length,
            num_path,
            rwr_rate,
            rate,
            num_e,
            num_n,
        );
        Ok((
            a.into_pyarray_bound(py),
            b.into_pyarray_bound(py),
            c.into_pyarray_bound(py),
            d.into_pyarray_bound(py),
            e.into_pyarray_bound(py),
        ))
    }

    fn sample_by_restart_random_walk<'py>(
        &self,
        py: Python<'py>,
        nodes: PyReadonlyArray1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        num_n: usize,
    ) -> PyResult<(Bound<'py, PyArray1<usize>>, Bound<'py, PyArray2<usize>>)> {
        let (a, b) = self.hg.sample_by_restart_random_walk(
            &nodes.as_array(),
            length,
            num_path,
            rwr_rate,
            num_n,
        );
        Ok((a.into_pyarray_bound(py), b.into_pyarray_bound(py)))
    }
}
/// A Python module implemented in Rust.
#[pymodule]
fn fish_utils(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyGraph>()?;
    m.add_class::<PyHyperGraph>()?;
    Ok(())
}
